# -*- coding: utf-8 -*-  
'''
transformer评价函数

@author: luoyi
Created on 2021年3月22日
'''
import tensorflow as tf

import utils.conf as conf
import utils.logger_factory as logf


#    预测准确率(T / P)
class TransformerMetric(tf.keras.metrics.Metric):
    
    def __init__(self, 
                 name='TransformerMetric',
                 batch_size=conf.DATASET.get_batch_size(),
                 sentence_maxlen=conf.TRANSFORMER.get_sentence_maxlen() + 1,
                 **kwargs):
        super(TransformerMetric, self).__init__(name=name, **kwargs)
        
        self._batch_size = batch_size
        self._sentence_maxlen = sentence_maxlen
        
        self.acc = self.add_weight('metric_acc', initializer='zero', dtype=tf.float64)
        pass
    
    def update_state(self, y_true, y_pred, sample_weight=None):
        '''
            @param y_pred: Tensor(batch_size, sentence_maxlen + 2, dict_size)
            @param y_pred: Tensor(batch_size, sentence_maxlen + 2)
        '''
        #    取y_pred每个sentence_maxlen的最大索引，并通过y_true去掉<PAD>部分
        y_pred = tf.math.argmax(y_pred, axis=-1)                        #    Tensor(batch_size, sentence_maxlen + 2)
        #    取y_true > 0的部分（实际需要预测的词）
        y_pred = tf.gather_nd(y_pred, indices=tf.where(y_true > 0))
        y_true = tf.gather_nd(y_true, indices=tf.where(y_true > 0))
        y_true = tf.cast(y_true, dtype=tf.int64)
        
        #    比较idx_pred与y_true的结果
        eq_res = tf.equal(y_pred, y_true)
        
        T = tf.math.count_nonzero(eq_res)
        P = tf.math.count_nonzero(y_true)
        p_acc = T / P
        self.acc.assign(p_acc)
        tf.print('y_true:', y_true, output_stream=logf.get_logger_filepath('transformer_metrics'))
        tf.print('y_pred:', y_pred, output_stream=logf.get_logger_filepath('transformer_metrics'))
        tf.print('acc:', p_acc, ' T:', T, ' P:', P, output_stream=logf.get_logger_filepath('transformer_metrics'))
        
        pass

    def reset_states(self):
        self.acc.assign(0.)
    def result(self):
        return self.acc
    
    pass